import numpy as np
import datetime
import tool


def calculate_weekday_totals(data):
    dates = np.array([datetime.datetime.strptime(d, "%Y-%m-%d") for d in data["date"]])
    weekdays = np.array([d.weekday() for d in dates])
    weekday_totals = np.zeros(7, dtype="f8")  # 7 天，每天一个和
    for i in range(7):
        mask = weekdays == i
        weekday_totals[i] = np.sum(data["total_quantity"][mask])
    result_dtype = [("weekday", "U10"), ("total_quantity", "f8")]
    weekdays_names = [
        "Monday",
        "Tuesday",
        "Wednesday",
        "Thursday",
        "Friday",
        "Saturday",
        "Sunday",
    ]
    result = np.array(
        [(weekdays_names[i], weekday_totals[i]) for i in range(7)], dtype=result_dtype
    )
    return result


if __name__ == "__main__":
    data = tool.get_np("total_daily_quantity.csv")
    result = calculate_weekday_totals(data)
    tool.get_csv(result, "cal_weekday_total.csv")
